package com.etl.component.file;

import cn.hutool.core.io.FileUtil;
import com.google.common.collect.Lists;
import com.component.api.constants.ComponentCodeConstants;
import com.component.api.constants.ComponentConstants;
import com.component.api.excepiton.ComponentException;
import com.component.api.model.PortData;
import com.component.api.model.ProcessResult;
import com.component.api.model.data.SetData;
import com.component.api.model.param.ParamPair;
import com.etl.component.common.AbstractFlinkComponent;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.Utils;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.operators.MapOperator;
import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.types.Row;

import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * 描述：
 * 文件写入
 *  太难了 flink还不太完善
 * @author xianggj
 * @Date 2021/11/4 15:13
 **/
public class FileWriteComponent extends AbstractFlinkComponent {
    @Override
    public ProcessResult process(List<PortData> datas, List<ParamPair> paramPairs) {
        Map<String, Object> params = paramToObject(paramPairs);
        String path = (String)params.get("path");
        //只支持 text 和csv
        String type = (String)params.get("type");
        //是否需要表头
        Boolean insertHeader = (Boolean) params.getOrDefault("header", true);
        SetData input = (SetData)datas.get(0).getValue();
        DataSet<Row> result = (DataSet<Row>) input.getDataset();
        //创建文件
        FileUtil.touch(path);

        if (ComponentConstants.CSV.equalsIgnoreCase(type)){
            MyCsvTableSink csvTableSink = new MyCsvTableSink(path, ",", 1,
                    FileSystem.WriteMode.OVERWRITE);
            csvTableSink.setEnableHeader(insertHeader);
            csvTableSink.consumeDataSet(result);
        } else {
            if (insertHeader){
                //插入表头
                result = insertHeader(result);
            }
            result.writeAsText(path, FileSystem.WriteMode.OVERWRITE)
                    .setParallelism(1);
        }
        try {
            ExecutionEnvironment executionEnvironment = result.getExecutionEnvironment();
            executionEnvironment.execute("执行文件写入操作");
        } catch (Exception e) {
            e.printStackTrace();
            throw new ComponentException("文件写入出现错误：" , e );
        }
        ProcessResult processResult = new ProcessResult();
        processResult.addAbstractTableData("output", input);
        return processResult;
    }

    /**
     * 插入表头
     * 这是个耗时操作 需要时在做添加
     * @param result
     * @return
     */
    private DataSet<Row> insertHeader(DataSet<Row> result) {
        RowTypeInfo rowTypeInfo = (RowTypeInfo)result.getType();
        String[] fieldNames = rowTypeInfo.getFieldNames();
        BasicTypeInfo[] basicTypeInfos = new BasicTypeInfo[fieldNames.length];
        Arrays.fill(basicTypeInfos, BasicTypeInfo.STRING_TYPE_INFO);
        RowTypeInfo newRowTypeInfo = new RowTypeInfo(basicTypeInfos, fieldNames);
        //转换原dataset为字符串
        result = toDatasetForStr(result, newRowTypeInfo);
        ExecutionEnvironment env = result.getExecutionEnvironment();
        DataSource<Row> headerDataSet = (DataSource<Row>) env.fromCollection(Arrays.asList(Row.of(fieldNames)), newRowTypeInfo);
        result = headerDataSet.union(result);
        return result;
    }

    private DataSet<Row> toDatasetForStr(DataSet<Row> result, RowTypeInfo newRowTypeInfo) {

        String callLocation = Utils.getCallLocationName();
        result = new MapOperator<Row, Row>(result, newRowTypeInfo, e ->{
            for (int i = 0; i < e.getArity(); i++) {
                e.setField(i, String.valueOf(e.getField(i)));
            }
            return e;
        }, callLocation);
        return result;
    }

    @Override
    public List<PortData> previewOutputColumns(List<PortData> datas, List<ParamPair> paramPairs) {
        return Lists.newArrayList();
    }

    @Override
    public String getRule() {
        return ComponentCodeConstants.FILE_OUTPUT;
    }
}
